from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
from sentence_transformers import SentenceTransformer

# 搜索参数
search_params = {"metric_type": "COSINE", "params": {"ef": 100}}
# 指定本地模型路径
model_path = r'E:\repository\models\paraphrase-multilingual-MiniLM-L12-v2'  # 替换为你的本地路径

# 加载本地模型
model = SentenceTransformer(model_path)
# 将查询文本转为向量
query_text = "简史"
query_embedding = model.encode([query_text])[0].tolist()

collection_name = "documents"
connections.connect("default", host="localhost", port="19530")
collection = Collection(name=collection_name)

# 执行搜索
results = collection.search(
    data=[query_embedding],           # 查询向量
    anns_field="embedding",           # 搜索的向量字段
    param=search_params,              # 搜索参数
    limit=5,                          # 返回Top5结果
    expr='category == "technology"',  # 过滤条件
    output_fields=["text", "category"] # 返回的字段
)

# 处理结果
for hits in results:
    for hit in hits:
        print(f"ID: {hit.id}, 相似度: {hit.score}")
        print(f"文本: {hit.entity.get('text')}")
        print(f"类别: {hit.entity.get('category')}")
        print("---")